
#ifndef PROTOCOLS_UDP_UDPAPP_ADAPTER_H_
#define PROTOCOLS_UDP_UDPAPP_ADAPTER_H_

#include <string>
#include <vector>
#include <map>
#include <unordered_map>
#include <memory>
#include <queue>


#include "ai_sim.h"
#include "inet/applications/udpapp/UdpBasicApp.h"
#include "RoceConnectionBase.h"
#include "inet/common/packet/Packet.h"

#include "messages/messages_m.h"
#include "RoceSenderManager.h"
#include "RoceReceiverManager.h"

using namespace inet;

// using namespace omnetpp;
namespace ai_sim {

const uint16_t UDP_MAX_MESSAGE_SIZE = 65535;  // bytes
const uint16_t RDMA_PAYLAOD_SIZE = 1000;  // bytes
const double QP_MINIMUM_RATE = 100; // Mbps
const double AI_RATE = 100; // 100Mbps
const double HAI_RATE = 400; // 400Mbps
const double LINE_RATE = 40000; // 40Gbps

class DcqcnRoceConnection : public RoceConnectionBase,public omnetpp::cSimpleModule {
public:
   DcqcnRoceConnection(const inet::L3Address& src_addr, const inet::L3Address& dst_addr,
        int src_port, int dst_port, RoceType type);
   ~DcqcnRoceConnection();

   /*************************
    * From RoceConnectionBase
    *************************/
   void HandlePacketFromHigherLayer(const FlowMsg& f_msg) override;
   void HandlePacketFromLowerLayer(inet::Packet& packet) override;
   inet::Packet* GetNextDataPacket(simtime_t* time) override;
   inet::Packet* GetNextControlPacket() override;
   FlowMsg* GetReceivedFlowMessage() override;

   /*****************************
    * extends from cSimpleModule
    *****************************/
   void handleMessage(cMessage *msg) override;
private:
   void scheduleTimeout(cMessage *msg, omnetpp::simtime_t timeout);
   
   void handleAck(uint32_t seqNum);
   void handleNack(uint32_t seqNum);
   /************
    * DCQCN
    ************/
   void handleCnp();
   void updateAlpha();
   void rateInc();
   void rateDecrease();
   void processTimer();

   //
   uint32_t getTransDelayNs(uint32_t sizeInBytes);
   inet::Packet *wrapPacket(uint32_t dataSize, 
         CLUSTER_MSGS kind, 
         uint32_t seqNum, 
         bool isLastSegment, 
         const Flow &flow);
   //
private: // data member
   struct {
      double targetRate;
      double curRate;
      double maxRate; // Mbps
      double minRate; // Mbps
      double alpha;
      bool alphaCnpArrived;
      bool firstCnpArrived;
      bool decreaseCnpArrived;
      // timer for DCQCN
      simtime_t alphaUpdateInter;
      simtime_t rateIncInter;
      double additiveRate;
      double hyperIncRate;
      uint32_t rateIncStage;
      double g;
   } dcqcn;

   struct { // queue pair metadata
      int srcPort;
      int dstPort;
      inet::L3Address srcAddr;
      inet::L3Address dstAddr;
      RoceType type;
   } qp;

   struct {
      inet::simtime_t cnpGap;
      inet::simtime_t lastCnpSentTime;
   } cnpSendInfo;

   struct {
      omnetpp::simtime_t dataGap;
      omnetpp::simtime_t lastSentTime;
   } dataSendInfo;

   struct {
      cMessage * alphaTimer;
      cMessage * rateIncTimer;
   } timers;

   std::queue<inet::Packet *> ctrlQueue;
   RoceReceiverManager receiverMgr;
   RoceSenderManager senderMgr{ 1000 };
}; // DCQCN_CONNECTION

}  // namespace ai_sim

#endif  // PROTOCOLS_UDP_UDPAPP_ADAPTER_H_
